#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <windows.h>
#include <direct.h>
#include <time.h>

#define BUFFER_SIZE 4*1024*1024  // 4MB buffer
#define PROGRESS_BAR_WIDTH 50
#define LOG_FILE "speed.txt"
#define PATH_MAX 4096
#define LOG_INTERVAL 10*1024*1024
#define MIN_ELAPSED_TIME 0.001
#define ENABLE_FSYNC 0          // 是否强制刷新到磁盘 (1=启用, 0=禁用)
#define SKIP_WARMUP_BLOCKS 2    // 跳过前N个块的速度计算

// 计算时间差（秒）
double time_diff(LARGE_INTEGER *start, LARGE_INTEGER *end, LARGE_INTEGER *frequency) {
    return (double)(end->QuadPart - start->QuadPart) / frequency->QuadPart;
}

// 显示进度条
void show_progress(double progress, double speed, double elapsed) {
    int bar_width = PROGRESS_BAR_WIDTH;
    int pos = (int)(bar_width * progress);
    
    printf("\r[");
    for (int i = 0; i < bar_width; ++i) {
        if (i < pos) printf("=");
        else if (i == pos) printf(">");
        else printf(" ");
    }
    printf("] %.2f%%  %.2f MB/s  ETA: ", progress * 100.0, speed);
    
    if (speed > 0) {
        long remaining = (long)((1.0 - progress) * elapsed / progress);
        if (remaining < 60) {
            printf("%ld sec", remaining);
        } else if (remaining < 3600) {
            printf("%ld min %ld sec", remaining / 60, remaining % 60);
        } else {
            printf("%ld hr %ld min", remaining / 3600, (remaining % 3600) / 60);
        }
    } else {
        printf("Calculating...");
    }
    
    fflush(stdout);
}

// 记录速度数据
void log_speed(FILE *log_file, double speed, int block_num) {
    if (log_file) {
        fprintf(log_file, "%.6f\n",speed);

    }
}

// 计算目录大小
__int64 calculate_directory_size(const char *path) {
    __int64 total_size = 0;
    char search_path[PATH_MAX];
    WIN32_FIND_DATAA find_data;
    HANDLE hFind;
    
    // 构建搜索路径
    sprintf_s(search_path, PATH_MAX, "%s\\*", path);
    
    // 查找第一个文件/目录
    hFind = FindFirstFileA(search_path, &find_data);
    if (hFind == INVALID_HANDLE_VALUE) {
        printf("Failed to open directory: %s\n", path);
        return 0;
    }
    
    do {
        // 跳过当前和上级目录
        if (strcmp(find_data.cFileName, ".") == 0 || strcmp(find_data.cFileName, "..") == 0)
            continue;
        
        // 构建子路径
        sprintf_s(search_path, PATH_MAX, "%s\\%s", path, find_data.cFileName);
        
        if (find_data.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
            // 递归计算子目录大小
            total_size += calculate_directory_size(search_path);
        } else {
            // 计算文件大小
            __int64 file_size = ((__int64)find_data.nFileSizeHigh << 32) | find_data.nFileSizeLow;
            total_size += file_size;
        }
    } while (FindNextFileA(hFind, &find_data) != 0);
    
    FindClose(hFind);
    return total_size;
}

// 复制单个文件
int copy_file(const char *src_path, const char *dst_path, __int64 file_size, 
              __int64 *total_bytes_copied, __int64 total_size, LARGE_INTEGER start_time, 
              FILE *log_file, __int64 *bytes_since_last_log, LARGE_INTEGER *block_start_time,
              double *moving_average_speed, int *block_counter, LARGE_INTEGER *frequency) {
    HANDLE src_handle, dst_handle;
    char *buffer = (char*)malloc(BUFFER_SIZE);
    if (!buffer) {
        printf("Memory allocation failed\n");
        return 0;
    }
    
    // 打开源文件
    src_handle = CreateFileA(src_path, GENERIC_READ, FILE_SHARE_READ, NULL, 
                             OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (src_handle == INVALID_HANDLE_VALUE) {
        printf("Failed to open source file: %s\n", src_path);
        free(buffer);
        return 0;
    }
    
    // 创建目标文件
    dst_handle = CreateFileA(dst_path, GENERIC_WRITE, 0, NULL, 
                             CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    if (dst_handle == INVALID_HANDLE_VALUE) {
        printf("Failed to create destination file: %s\n", dst_path);
        free(buffer);
        CloseHandle(src_handle);
        return 0;
    }
    
    DWORD bytes_read, bytes_written;
    __int64 file_bytes_copied = 0;
    
    while (ReadFile(src_handle, buffer, BUFFER_SIZE, &bytes_read, NULL) && bytes_read > 0) {
        if (!WriteFile(dst_handle, buffer, bytes_read, &bytes_written, NULL) || bytes_written != bytes_read) {
            printf("Failed to write to destination file: %s\n", dst_path);
            free(buffer);
            CloseHandle(src_handle);
            CloseHandle(dst_handle);
            return 0;
        }
        
        file_bytes_copied += bytes_read;
        *total_bytes_copied += bytes_read;
        *bytes_since_last_log += bytes_read;
        
        // 检查是否需要记录速度
        if (*bytes_since_last_log >= LOG_INTERVAL) {
            LARGE_INTEGER current_time;
            QueryPerformanceCounter(&current_time);
            double block_elapsed = time_diff(block_start_time, &current_time, frequency);
            
            // 强制刷新到磁盘（如果启用）
            if (ENABLE_FSYNC) {
                FlushFileBuffers(dst_handle);
            }
            
            // 重新获取时间，确保包含刷新时间
            QueryPerformanceCounter(&current_time);
            block_elapsed = time_diff(block_start_time, &current_time, frequency);
            
            // 跳过预热阶段的高速块
            (*block_counter)++;
            if (*block_counter <= SKIP_WARMUP_BLOCKS) {
                *bytes_since_last_log = 0;
                *block_start_time = current_time;
                continue;
            }
            
            // 确保时间间隔安全
            if (block_elapsed < MIN_ELAPSED_TIME) {
                block_elapsed = MIN_ELAPSED_TIME;
            }
            
            double block_speed = (double)*bytes_since_last_log / (1024 * 1024) / block_elapsed;
            // 记录速度
            log_speed(log_file, block_speed, *block_counter);
            
            // 更新移动平均速度
            if (*moving_average_speed == 0) {
                *moving_average_speed = block_speed;
            } else {
                *moving_average_speed = 0.9 * *moving_average_speed + 0.1 * block_speed;
            }
            
            *bytes_since_last_log = 0;
            *block_start_time = current_time;
            
            // 更新进度显示
            double elapsed_seconds = time_diff(&start_time, &current_time, frequency);
            double progress = (double)*total_bytes_copied / total_size;
            show_progress(progress, *moving_average_speed, elapsed_seconds);
        }
    }
    
    free(buffer);
    CloseHandle(src_handle);
    CloseHandle(dst_handle);
    return 1;
}

// 递归复制目录
int copy_directory(const char *src_path, const char *dst_path, __int64 total_size, 
                   __int64 *total_bytes_copied, LARGE_INTEGER start_time, FILE *log_file,
                   __int64 *bytes_since_last_log, LARGE_INTEGER *block_start_time,
                   double *moving_average_speed, int *block_counter, LARGE_INTEGER *frequency) {
    char search_path[PATH_MAX];
    WIN32_FIND_DATAA find_data;
    HANDLE hFind;
    
    // 创建目标目录
    if (_mkdir(dst_path) == -1) {
        printf("Failed to create directory: %s\n", dst_path);
        return 0;
    }
    
    // 构建搜索路径
    sprintf_s(search_path, PATH_MAX, "%s\\*", src_path);
    
    // 查找第一个文件/目录
    hFind = FindFirstFileA(search_path, &find_data);
    if (hFind == INVALID_HANDLE_VALUE) {
        printf("Failed to open directory: %s\n", src_path);
        return 0;
    }
    
    do {
        // 跳过当前和上级目录
        if (strcmp(find_data.cFileName, ".") == 0 || strcmp(find_data.cFileName, "..") == 0)
            continue;
        
        char src_subpath[PATH_MAX], dst_subpath[PATH_MAX];
        
        // 构建源和目标子路径
        sprintf_s(src_subpath, PATH_MAX, "%s\\%s", src_path, find_data.cFileName);
        sprintf_s(dst_subpath, PATH_MAX, "%s\\%s", dst_path, find_data.cFileName);
        
        if (find_data.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
            // 递归复制子目录
            if (!copy_directory(src_subpath, dst_subpath, total_size, 
                               total_bytes_copied, start_time, log_file, 
                               bytes_since_last_log, block_start_time,
                               moving_average_speed, block_counter, frequency)) {
                FindClose(hFind);
                return 0;
            }
        } else {
            // 计算文件大小
            __int64 file_size = ((__int64)find_data.nFileSizeHigh << 32) | find_data.nFileSizeLow;
            
            // 复制文件
            if (!copy_file(src_subpath, dst_subpath, file_size, 
                          total_bytes_copied, total_size, start_time, 
                          log_file, bytes_since_last_log, block_start_time,
                          moving_average_speed, block_counter, frequency)) {
                FindClose(hFind);
                return 0;
            }
        }
    } while (FindNextFileA(hFind, &find_data) != 0);
    
    FindClose(hFind);
    return 1;
}

int main(int argc, char *argv[]) {
    if (argc != 3) {
        printf("Usage: %s source_directory destination_directory\n", argv[0]);
        return 1;
    }
    
    const char *source_dir = argv[1];
    const char *dest_dir = argv[2];
    
    // 检查源目录是否存在
    DWORD file_attr = GetFileAttributesA(source_dir);
    if (file_attr == INVALID_FILE_ATTRIBUTES || !(file_attr & FILE_ATTRIBUTE_DIRECTORY)) {
        printf("Error: Source path '%s' does not exist or is not a directory\n", source_dir);
        return 1;
    }
    
    // 检查目标目录是否存在
    file_attr = GetFileAttributesA(dest_dir);
    if (file_attr != INVALID_FILE_ATTRIBUTES) {
        if (!(file_attr & FILE_ATTRIBUTE_DIRECTORY)) {
            printf("Error: Destination path '%s' exists but is not a directory\n", dest_dir);
            return 1;
        }
        
        char response[10];
        printf("Destination directory '%s' already exists.Please change dest dir.", dest_dir);
        return 1;
    }
    
    // 获取高精度计时器频率
    LARGE_INTEGER frequency;
    if (!QueryPerformanceFrequency(&frequency)) {
        printf("High-resolution performance counter not supported\n");
        return 1;
    }
    
    printf("Calculating directory size...\n");
    __int64 total_size = calculate_directory_size(source_dir);
    if (total_size == 0) {
        printf("Source directory is empty\n");
        return 0;
    }
    
    FILE *log_file = fopen(LOG_FILE, "w");
    if (!log_file) {
        printf("Failed to create speed log file\n");
        return 1;
    }
    
    printf("Starting directory copy: %s -> %s\n", source_dir, dest_dir);
    printf("Total size: %.2f MB\n", (double)total_size / (1024 * 1024));
    
    __int64 total_bytes_copied = 0;
    __int64 bytes_since_last_log = 0;
    LARGE_INTEGER start_time, block_start_time;
    QueryPerformanceCounter(&start_time);
    block_start_time = start_time;
    double moving_average_speed = 0;
    int block_counter = 0;
    
    if (!copy_directory(source_dir, dest_dir, total_size, &total_bytes_copied, 
                       start_time, log_file, &bytes_since_last_log, &block_start_time,
                       &moving_average_speed, &block_counter, &frequency)) {
        fclose(log_file);
        return 1;
    }
    
    // 处理剩余数据
    if (bytes_since_last_log > 0 && block_counter > SKIP_WARMUP_BLOCKS) {
        LARGE_INTEGER end_time;
        QueryPerformanceCounter(&end_time);
        double block_elapsed = time_diff(&block_start_time, &end_time, &frequency);
        
        if (block_elapsed < MIN_ELAPSED_TIME) {
            block_elapsed = MIN_ELAPSED_TIME;
        }
        
        double block_speed = (double)bytes_since_last_log / (1024 * 1024) / block_elapsed;
        printf("Final block: %.2f MB/s (%.6f seconds)\n", block_speed, block_elapsed);
        log_speed(log_file, block_speed, block_counter + 1);
        
        if (moving_average_speed == 0) {
            moving_average_speed = block_speed;
        } else {
            moving_average_speed = 0.9 * moving_average_speed + 0.1 * block_speed;
        }
    }
    fflush(log_file);
    fclose(log_file);
    
    // 计算最终统计信息
    LARGE_INTEGER end_time;
    QueryPerformanceCounter(&end_time);
    double total_elapsed = time_diff(&start_time, &end_time, &frequency);
    double average_speed = (double)total_size / (1024 * 1024) / total_elapsed;
    
    printf("\r[");
    for (int i = 0; i < PROGRESS_BAR_WIDTH; ++i) printf("=");
    printf("] 100.00%%  %.2f MB/s  Total time: %.2f seconds\n", 
           average_speed, total_elapsed);
    
    printf("Directory copy completed: %s -> %s\n", source_dir, dest_dir);
    printf("Speed records saved to: %s\n", LOG_FILE);
    
    return 0;
}